from contextlib import contextmanager, AbstractContextManager
from os import PathLike
from pathlib import Path
from typing import AnyStr, Type, Dict

import lmdb

from storage.kvstore.base import AbstractKVStore
from storage.mutex import Mutex, MutexProtocol


class LMDBKVStore(AbstractKVStore):
    def __init__(self,
                 path: AnyStr | PathLike,
                 *,
                 lock_type: Type[MutexProtocol] = Mutex,
                 map_size=100 * 1024 * 1024,
                 sync=True):
        self._path = Path(path)
        self._data = lmdb.open(str(self._path), map_size=map_size, sync=sync)
        self._lock_type = lock_type
        self._lock_path = self._path.resolve().with_suffix('.lock')

    def close(self):
        self._data.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def get(self, k: bytes, default: bytes = None) -> bytes:
        with self._data.begin() as cursor:
            return cursor.get(k)

    def put(self, k: bytes, v: bytes):
        with self._data.begin(write=True) as cursor:
            cursor.put(k, v)

    def update(self, kv_pairs: Dict[bytes, bytes]):
        with self._data.begin(write=True) as txn:
            cursor = txn.cursor()
            cursor.putmulti(kv_pairs.items())

    def __delitem__(self, k: bytes):
        with self._data.begin(write=True) as txn:
            txn.delete(k)

    def __getitem__(self, item: bytes):
        with self._data.begin() as txn:
            if (ret := txn.get(item)) is not None:
                return ret
            raise KeyError(item)

    def __setitem__(self, key: bytes, value: bytes):
        with self._data.begin(write=True) as txn:
            txn.put(key, value)

    @contextmanager
    def tx(self) -> AbstractContextManager:
        with self._lock_type(str(self._lock_path)):
            yield self
